"""
简化版垃圾邮件检测器
"""
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer     # TF-IDF向量化
from sklearn.naive_bayes import MultinomialNB                   # 朴素贝叶斯分类器   
from sklearn.model_selection import train_test_split            # 数据集划分
from sklearn.metrics import classification_report, accuracy_score   # 评估指标
import re

def preprocess_text(text):
    """简单的文本预处理"""
    text = text.lower()
    text = re.sub(r'[^a-zA-Z\s]', '', text)
    return text

def main():
    # 创建示例数据
    data = {
        'text': [
            "URGENT! You have won $1000000! Claim now!",
            "FREE VIAGRA! Buy now! Limited time offer!",
            "Hi John, how are you doing today?",
            "Meeting scheduled for 3 PM in conference room A.",
            "CLICK HERE FOR FREE MONEY! Limited time only!",
            "Your package has been delivered successfully.",
            "CASINO BONUS! Play now and win! Free spins!",
            "Thank you for your presentation today.",
            "MIRACLE CURE! Doctors hate this one trick!",
            "Flight confirmation: Your flight to New York departs at 8 AM."
        ],
        'label': ['spam', 'spam', 'ham', 'ham', 'spam', 'ham', 'spam', 'ham', 'spam', 'ham']
    }
    
    df = pd.DataFrame(data)
    
    # 预处理文本
    df['processed_text'] = df['text'].apply(preprocess_text)
    
    # 创建特征
    vectorizer = TfidfVectorizer(max_features=1000, stop_words='english')
    X = vectorizer.fit_transform(df['processed_text'])
    y = df['label'].map({'ham': 0, 'spam': 1})
    
    # 分割数据
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
    
    # 训练模型
    model = MultinomialNB()
    model.fit(X_train, y_train)
    
    # 预测和评估
    y_pred = model.predict(X_test)
    print(f"准确率: {accuracy_score(y_test, y_pred):.4f}")
    print("\n分类报告:")
    print(classification_report(y_test, y_pred, target_names=['Ham', 'Spam']))
    
    # 测试新邮件
    test_email = "FREE GIFT CARD! Claim your $500 reward now!"
    processed_test = preprocess_text(test_email)
    test_vector = vectorizer.transform([processed_test])
    prediction = model.predict(test_vector)[0]
    probability = model.predict_proba(test_vector)[0]
    
    print(f"\n测试邮件: {test_email}")
    print(f"预测结果: {'Spam' if prediction == 1 else 'Ham'}")
    print(f"垃圾邮件概率: {probability[1]:.4f}")

if __name__ == "__main__":
    main()